﻿#pragma once
#include <iostream>
#include <cassert>
using namespace std;

enum COLOUR
{
	RED,
	BLACK
};

template<class V>
struct RBTNode
{
	V _val;
	RBTNode<V>* _left = nullptr;
	RBTNode<V>* _right = nullptr;
	RBTNode<V>* _parent = nullptr;
	COLOUR _color;

	RBTNode(V val = V())
		:_val(val)
	{}
};


//实现红黑树的迭代器:为双向迭代器
template<class V, class Ptr, class Ref>
struct iteratorForRBTree
{
	typedef RBTNode<V> Node;
	typedef iteratorForRBTree<V, Ptr, Ref> Self;
	iteratorForRBTree(Node* node,Node* root)
		:_node(node)
		,_root(root)
	{}
	//重载++,--,*,->,!=,==

	//前置++
	Self& operator++()
	{
		//红黑树走中序遍历
		if (_root)
		{
			if (_node->_right)
			{
				//如果当前结点的右树不为空，走到右树的最左结点
				_node = _node->_right;
				while (_node->_left)
				{
					_node = _node->_left;
				}
			}
			else
			{
				//当前节点的右树为空，向上走
				Node* parent = _node->_parent;
				//是左孩子，还是右孩子呢？
				while (parent && parent->_right == _node)
				{
					_node = parent;
					parent = parent->_parent;
				}

				//出循环有两种可能：parent为空，或是左孩子

				//parent为空，说明整棵树已经遍历完
				if (parent == nullptr)
					_node = nullptr;
				//_node是parent的左孩子
				else
					_node = parent;

				//其实可以直接写成_node = parent		
			}
		}
		
		return *this;
	}

	//后置++
	Self operator++(int)
	{
		Self tmp = *this;//调用默认生成的拷贝构造
		this->operator++();
		return tmp;
	}

	//前置--
	Self& operator--()
	{
		if (_root)
		{
			//_node为空时，说明整棵树走完了
			if (_node == nullptr)
			{
				_node = _root->_right;
				while (_node->_right)
					_node = _node->_right;
			}
			else
			{
				//_node不为空时,考虑左子树为空或不为空
				if (_node->_left)
				{
					_node = _node->_left;
					while (_node->_right)
						_node = _node->_right;
				}
				else
				{
					Node* parent = _node->_parent;
					//是左孩子，还是右孩子呢？
					while (parent && parent->_left == _node)
					{
						_node = parent;
						parent = parent->_parent;
					}

					//出循环，两种可能
					if (parent == nullptr)
						assert(false);//非法操作
					else
						_node = parent;
				}
			}
		}
		return *this;
	}

	//后置--
	Self operator--(int)
	{
		Self tmp = *this;
		this->operator--();
		return tmp;
	}

	//当函数的返回值是引用时，无需创建临时变量，直接返回某个已存在的对象
	Ref operator*()
	{
		return _node->_val;
	}

	Ptr operator->()
	{
		return &_node->_val;
	}

	bool operator!=(const Self& it)
	{
		return _node != it._node;
	}

	bool operator==(const Self& it)
	{
		return _node == it._node;
	}
private:
	Node* _node;
	Node* _root;
};


//使用红黑树来封装map和set
//由于不确定到底是k还是kv结构,所以使用K代表键值，V代表实际数据
//用一个仿函数kget获取K值，个人感觉这个仿函数与K有点重合

template<class K, class V, class kGet>
class RBTree
{
public:
	typedef RBTNode<V> Node;
	typedef iteratorForRBTree<V, V*, V&> Iterator;
	typedef iteratorForRBTree<V, const V*, const V&> Const_Iterator;

	RBTree() = default;
	
	~RBTree()
	{
		destroy(_root);
	}

	void destroy(Node* root)
	{
		if (root == nullptr)
			return;
		destroy(root->_left);
		destroy(root->_right);
		delete root;
	}

	RBTree(const RBTree<K,V,kGet>& t)
	{
		auto it = t.begin();
		while (it != t.end())
		{
			this->insert(*it);
			it++;
		}
	}

	RBTree<K, V, kGet>& operator=(const RBTree<K, V, kGet>& t)
	{
		RBTree<K, V, kGet> tmp = t;
		swap(this->_root, tmp._root);
		return *this;
	}


	pair<Iterator, bool> insert(const V& val)
	{
		Node* ptr = new Node(val);
		ptr->_color = RED;
		if (_root == nullptr)
		{
			_root = ptr;
			_root->_color = BLACK;
		}
		else
		{
			//遍历整棵二叉搜索树，找到应处位置
			Node* cur = _root;
			Node* pre = nullptr;
			while (cur)
			{
				pre = cur;
				if (kGet()(ptr->_val) < kGet()(cur->_val))
				{
					cur = cur->_left;
				}
				else if (kGet()(ptr->_val) > kGet()(cur->_val))
				{
					cur = cur->_right;
				}
				else//暂且不处理相等的情况
				{
					delete ptr;
					return { {cur,_root},false };
				}

			}
			ptr->_parent = pre;
			if (kGet()(ptr->_val) < kGet()(pre->_val))
			{
				pre->_left = ptr;
			}
			else
			{
				pre->_right = ptr;
			}
		}

		Node* cur = ptr;
		Node* parent = ptr->_parent;
		//插入结点，初始化完成，开始变色+旋转

		//1.如果parent为空(根结点插入情况)或者parent为黑，不需要处理

		//2. 需要处理的是parent不为空并且parent为红的情况
		while (parent && parent->_color == RED)
		{
			Node* grandParent = parent->_parent;
			if (parent == grandParent->_left)
			{
				Node* uncle = grandParent->_right;
				//分三种情况，但两种处理来讨论

				//1.叔叔存在且为红，不断向上调整,此时仅变色即可
				//注意，不能写成循环的形式
				if (uncle && uncle->_color == RED)
				{
					parent->_color = uncle->_color = BLACK;
					grandParent->_color = RED;

					cur = grandParent;
					parent = cur->_parent;
				}

				//2.叔叔存在且为黑，或叔叔不存在，此时需变色+旋转
				//同样内部不能写成循环的形式
				else if (uncle == nullptr || uncle && uncle->_color == BLACK)
				{
					if (cur == parent->_left)
					{
						//单旋+变色
						RotateR(parent);

						parent->_color = BLACK;
						grandParent->_color = RED;
						break;
					}
					else if (cur == parent->_right)
					{
						//双旋+变色
						RotateL(cur);
						RotateR(cur);

						cur->_color = BLACK;
						grandParent->_color = RED;
						break;
					}
					else
					{
						assert(false);
					}
				}
				else
				{
					assert(false);
				}

			}
			else if (parent == grandParent->_right)
			{
				Node* uncle = grandParent->_left;
				//分三种情况，但两种处理来讨论

				//1.叔叔存在且为红，不断向上调整,此时仅变色即可
				//注意，不能写成循环的形式
				if (uncle && uncle->_color == RED)
				{
					parent->_color = uncle->_color = BLACK;
					grandParent->_color = RED;

					cur = grandParent;
					parent = cur->_parent;
				}

				//2.叔叔存在且为黑，或叔叔不存在，此时需变色+旋转
				//同样内部不能写成循环的形式
				else if (uncle == nullptr || uncle && uncle->_color == BLACK)
				{
					if (cur == parent->_right)
					{
						//单旋+变色
						RotateL(parent);

						parent->_color = BLACK;
						grandParent->_color = RED;
						break;
					}
					else if (cur == parent->_left)
					{
						//双旋+变色
						RotateR(cur);
						RotateL(cur);

						cur->_color = BLACK;
						grandParent->_color = RED;
						break;
					}
					else
					{
						assert(false);
					}
				}
				else
				{
					assert(false);
				}
			}
			else
			{
				assert(false);
			}
		}
		_root->_color = BLACK;
		return { {ptr,_root},true };
	}

	Iterator find(const K& key)
	{
		Node* cur = _root;
		while (cur)
		{
			if (kGet()(cur->_val) == key)
				return { cur,_root };
			else if (kGet()(cur->_val) < key)
				cur = cur->_right;
			else
				cur = cur->_left;
		}
		return { nullptr,_root };
	}


	//迭代器相关

	Iterator begin()
	{
		//返回中序遍历的第一个
		Node* cur = _root;
		while (cur && cur->_left)
			cur = cur->_left;
		return Iterator(cur,_root);
	}

	Const_Iterator begin() const
	{
		Node* cur = _root;
		while (cur && cur->_left)
			cur = cur->_left;
		return Const_Iterator(cur,_root);
	}

	Iterator end()
	{
		return Iterator(nullptr,_root);
	}

	Const_Iterator end() const
	{
		return Const_Iterator(nullptr,_root);
	}

private:
	Node* _root = nullptr;

	//左右旋转的函数设置为私有，仅内部调用，不允许外部调用
	//右单旋
	void RotateR(Node* son)
	{
		Node* parent = son->_parent;
		Node* sonR = son->_right;
		Node* gparent = parent->_parent;
		son->_right = parent;
		parent->_parent = son;
		parent->_left = sonR;
		if (sonR)
			sonR->_parent = parent;

		son->_parent = gparent;
		if (gparent == nullptr)
		{
			//说明parent为根结点，旋转后，son为根结点
			_root = son;
		}
		else
		{
			if (gparent->_left == parent)
				gparent->_left = son;
			else if (gparent->_right == parent)
				gparent->_right = son;
			else
				assert(false);
		}
	}

	//左单旋
	void RotateL(Node* son)
	{
		Node* parent = son->_parent;
		Node* gparent = parent->_parent;
		Node* sonL = son->_left;

		son->_left = parent;
		parent->_parent = son;
		parent->_right = sonL;
		if (sonL)
			sonL->_parent = parent;

		son->_parent = gparent;
		if (gparent == nullptr)
		{
			_root = son;
		}
		else
		{
			if (gparent->_left == parent)
			{
				gparent->_left = son;
			}
			else if (gparent->_right == parent)
			{
				gparent->_right = son;
			}
			else
				assert(false);
		}
	}
};

